import os
import logging
import requests
from tqdm import tqdm
import json
import time

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def download_file(url, local_path, headers=None):
    """下载单个文件"""
    try:
        if headers is None:
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
            }
        
        logger.info(f"开始下载: {url}")
        response = requests.get(url, stream=True, headers=headers)
        response.raise_for_status()  # 检查响应状态
        
        total_size = int(response.headers.get('content-length', 0))
        logger.info(f"文件大小: {total_size/1024/1024:.2f} MB")
        
        with open(local_path, 'wb') as f, tqdm(
            desc=os.path.basename(local_path),
            total=total_size,
            unit='iB',
            unit_scale=True,
            unit_divisor=1024,
        ) as pbar:
            for data in response.iter_content(chunk_size=8192):
                size = f.write(data)
                pbar.update(size)
        
        # 验证文件大小
        if os.path.getsize(local_path) != total_size:
            raise Exception(f"文件大小不匹配: 预期 {total_size}, 实际 {os.path.getsize(local_path)}")
            
        logger.info(f"下载完成: {local_path}")
        return True
        
    except Exception as e:
        logger.error(f"下载失败: {str(e)}")
        if os.path.exists(local_path):
            os.remove(local_path)
        return False

def get_model_files(model_id):
    """获取模型文件列表"""
    try:
        url = f"https://huggingface.co/{model_id}/raw/main/config.json"
        response = requests.get(url)
        response.raise_for_status()
        config = response.json()
        
        # 获取模型文件列表
        files_url = f"https://huggingface.co/{model_id}/tree/main"
        response = requests.get(files_url)
        response.raise_for_status()
        
        # 这里需要解析HTML来获取文件列表
        # 为了简单起见，我们直接返回已知的关键文件
        return [
            "config.json",
            "pytorch_model.bin",
            "tokenizer.json",
            "tokenizer_config.json",
            "special_tokens_map.json",
            "model.safetensors",
            "generation_config.json"
        ]
    except Exception as e:
        logger.error(f"获取文件列表失败: {str(e)}")
        return []

def main():
    # 模型信息
    model_id = "Qwen/Qwen3-Embedding-0.6B"
    local_dir = "./models/Qwen3-Embedding-0.6B"
    
    # 创建本地目录
    os.makedirs(local_dir, exist_ok=True)
    
    # 设置请求头
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
    }
    
    # 获取文件列表
    files_to_download = get_model_files(model_id)
    if not files_to_download:
        logger.error("无法获取文件列表，使用默认文件列表")
        files_to_download = [
            "config.json",
            "pytorch_model.bin",
            "tokenizer.json",
            "tokenizer_config.json",
            "special_tokens_map.json",
            "model.safetensors",
            "generation_config.json"
        ]
    
    base_url = f"https://huggingface.co/{model_id}/resolve/main/"
    
    # 下载文件
    success_count = 0
    for file in files_to_download:
        url = base_url + file
        local_path = os.path.join(local_dir, file)
        
        # 如果文件已存在且大小大于0，跳过下载
        if os.path.exists(local_path) and os.path.getsize(local_path) > 0:
            logger.info(f"文件已存在，跳过: {file}")
            success_count += 1
            continue
        
        # 尝试下载，最多重试3次
        for attempt in range(3):
            if download_file(url, local_path, headers):
                success_count += 1
                break
            logger.warning(f"第 {attempt + 1} 次下载失败，等待后重试...")
            time.sleep(5)
    
    # 检查下载结果
    if success_count == len(files_to_download):
        logger.info("所有文件下载成功！")
    else:
        logger.warning(f"部分文件下载失败: {success_count}/{len(files_to_download)}")
        
    # 列出下载的文件
    logger.info("\n下载的文件列表:")
    for file in os.listdir(local_dir):
        size = os.path.getsize(os.path.join(local_dir, file))
        logger.info(f"{file}: {size/1024/1024:.2f} MB")

if __name__ == "__main__":
    main() 